#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gmssl/sm2.h>
#include <gmssl/rand.h>
#include <gmssl/sm3.h>
#include <gmssl/error.h>
#include "../../include/crypto/sm2_keyex.h"

// 安全清零函数
static void secure_zero(void *ptr, size_t size)
{
    volatile unsigned char *p = ptr;
    while (size--)
        *p++ = 0;
}

// 生成SM2密钥对
int sm2_generate_key_pair(sm2_key_pair_t *key_pair)
{
    if (!key_pair)
        return -1;

    key_pair->has_private_key = 1;

    if (sm2_key_generate(&key_pair->key) != 1)
    {
        fprintf(stderr, "生成SM2密钥对失败\n");
        return -1;
    }

    return 0;
}

// 从私钥字节数组创建SM2密钥对
int sm2_create_key_pair_from_private_key(sm2_key_pair_t *key_pair,
                                         const unsigned char *private_key,
                                         size_t key_len)
{
    if (!key_pair || !private_key || key_len != 32)
    {
        return -1;
    }

    key_pair->has_private_key = 1;

    if (sm2_key_set_private_key(&key_pair->key, private_key) != 1)
    {
        fprintf(stderr, "设置SM2私钥失败\n");
        return -1;
    }

    return 0;
}

// 获取公钥坐标
int sm2_get_public_key_coords(const sm2_key_pair_t *key_pair,
                              unsigned char *x, size_t x_len,
                              unsigned char *y, size_t y_len)
{
    if (!key_pair || !x || !y || x_len < 32 || y_len < 32)
        return -1;

    // SM2_KEY 结构体中的 public_key 本身就是 SM2_POINT 类型
    // 直接从 SM2_POINT 结构中复制坐标
    memcpy(x, key_pair->key.public_key.x, 32);
    memcpy(y, key_pair->key.public_key.y, 32);

    return 0;
}

// SM3 KDF实现
static int sm3_kdf(const unsigned char *z, size_t z_len,
                   unsigned char *key, size_t key_len)
{
    SM3_CTX ctx;
    unsigned char digest[SM3_DIGEST_SIZE];
    unsigned int i, rounds;
    unsigned int counter = 1;
    size_t offset = 0;

    if (!z || !key || !key_len)
        return -1;

    rounds = (key_len + SM3_DIGEST_SIZE - 1) / SM3_DIGEST_SIZE;

    for (i = 0; i < rounds; i++)
    {
        sm3_init(&ctx);
        sm3_update(&ctx, z, z_len);

        // 添加计数器（大端序）
        unsigned char counter_bytes[4];
        counter_bytes[0] = (counter >> 24) & 0xFF;
        counter_bytes[1] = (counter >> 16) & 0xFF;
        counter_bytes[2] = (counter >> 8) & 0xFF;
        counter_bytes[3] = counter & 0xFF;

        sm3_update(&ctx, counter_bytes, sizeof(counter_bytes));
        sm3_finish(&ctx, digest);

        size_t copy_len = (key_len - offset) > SM3_DIGEST_SIZE ? SM3_DIGEST_SIZE : (key_len - offset);
        memcpy(key + offset, digest, copy_len);

        offset += copy_len;
        counter++;

        if (offset >= key_len)
            break;
    }

    // 安全清除上下文和临时摘要
    secure_zero(&ctx, sizeof(SM3_CTX));
    secure_zero(digest, sizeof(digest));

    return 0;
}

// 计算共享密钥
int sm2_derive_shared_key(const sm2_key_pair_t *local_key_pair,
                          const unsigned char *peer_public_x, size_t peer_x_len,
                          const unsigned char *peer_public_y, size_t peer_y_len,
                          unsigned char *shared_key, size_t shared_key_len)
{
    if (!local_key_pair || !peer_public_x || !peer_public_y || !shared_key ||
        peer_x_len < 32 || peer_y_len < 32 || !shared_key_len || !local_key_pair->has_private_key)
    {
        return -1;
    }

    int ret = -1;
    SM2_POINT peer_point;
    SM2_POINT shared_point;
    unsigned char z[64]; // 用于存储共享点的x和y坐标

    // 设置对方公钥点
    memcpy(peer_point.x, peer_public_x, 32);
    memcpy(peer_point.y, peer_public_y, 32);

    // 计算共享点
    if (sm2_do_ecdh(&local_key_pair->key, &peer_point, &shared_point) != 1)
    {
        fprintf(stderr, "SM2 ECDH密钥协商失败\n");
        goto cleanup;
    }

    // 将共享点的坐标拼接作为KDF输入
    memcpy(z, shared_point.x, 32);
    memcpy(z + 32, shared_point.y, 32);

    // 使用KDF生成共享密钥
    if (sm3_kdf(z, sizeof(z), shared_key, shared_key_len) != 0)
    {
        fprintf(stderr, "KDF派生共享密钥失败\n");
        goto cleanup;
    }

    ret = 0;

cleanup:
    // 安全清除敏感数据
    secure_zero(&shared_point, sizeof(shared_point));
    secure_zero(z, sizeof(z));

    return ret;
}

// 释放密钥对资源
void sm2_free_key_pair(sm2_key_pair_t *key_pair)
{
    if (key_pair)
    {
        // GmSSL的SM2_KEY不需要特别释放，只需要安全清零
        if (key_pair->has_private_key)
        {
            secure_zero(&key_pair->key, sizeof(SM2_KEY));
            key_pair->has_private_key = 0;
        }
    }
}

// SM2加密密钥
int sm2_encrypt_key(const unsigned char *plaintext, size_t plaintext_len,
                    const unsigned char *peer_public_x, size_t peer_x_len,
                    const unsigned char *peer_public_y, size_t peer_y_len,
                    unsigned char *ciphertext, size_t *ciphertext_len)
{
    if (!plaintext || !plaintext_len || !peer_public_x || !peer_public_y ||
        !ciphertext || !ciphertext_len || peer_x_len < 32 || peer_y_len < 32)
    {
        return -1;
    }

    // SM2加密后的密文格式: C1(65字节) || C3(32字节) || C2(明文长度)
    size_t expected_ciphertext_len = 65 + 32 + plaintext_len;

    if (*ciphertext_len < expected_ciphertext_len)
    {
        *ciphertext_len = expected_ciphertext_len;
        return -2; // 缓冲区太小
    }

    SM2_KEY peer_key;
    SM2_POINT point;

    // 构建对方公钥
    memcpy(point.x, peer_public_x, 32);
    memcpy(point.y, peer_public_y, 32);

    if (sm2_key_set_public_key(&peer_key, &point) != 1)
    {
        fprintf(stderr, "设置SM2公钥失败\n");
        return -1;
    }

    // 执行SM2加密
    size_t outlen = *ciphertext_len;
    if (sm2_encrypt(&peer_key, plaintext, plaintext_len, ciphertext, &outlen) != 1)
    {
        fprintf(stderr, "SM2加密失败\n");
        return -1;
    }

    *ciphertext_len = outlen;
    return 0;
}

// SM2解密密钥
int sm2_decrypt_key(const unsigned char *ciphertext, size_t ciphertext_len,
                    const sm2_key_pair_t *key_pair,
                    unsigned char *plaintext, size_t *plaintext_len)
{
    if (!ciphertext || !ciphertext_len || !key_pair || !plaintext ||
        !plaintext_len || !key_pair->has_private_key)
    {
        return -1;
    }

    // SM2密文的最小长度是C1(65字节) + C3(32字节) + 至少1字节的C2
    if (ciphertext_len < (65 + 32 + 1))
    {
        fprintf(stderr, "SM2密文长度无效\n");
        return -1;
    }

    // 执行SM2解密
    size_t outlen = *plaintext_len;
    if (sm2_decrypt(&key_pair->key, ciphertext, ciphertext_len, plaintext, &outlen) != 1)
    {
        fprintf(stderr, "SM2解密失败\n");
        return -1;
    }

    *plaintext_len = outlen;
    return 0;
}